/**
 * @file gvozdarev_3.c
 * @author Gvozdarev Sergey, gr016 MIPT
 * MIT license
 */
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include "mpi.h"


#define G_ELEM_NUM 1
#define G_FINISH   -1
#define G_MAIN_RANK 0

typedef float real_t;
#define MPI_REAL MPI_FLOAT
//
// сходится всегда
// апроксимация - h*h << tau 
const real_t h   = 0.01;
const real_t tau = 0.01 * 0.01 / 100.0;


real_t X;
real_t T;

real_t f_t( real_t t);
real_t f_x( real_t x);
real_t u( int i, real_t ** U);
real_t right_u( int i, real_t ** U);

void optim_distrib( int comm_size, int data_size, int * displs, int * counts);


int main( int argc, char ** argv)  
{
	int comm_rank,comm_size;
    MPI_Status *  status = NULL;

    // check arg
    if ( argc != 3)
    {
        printf( "X and T must be defined\n");
        return 1;
    }

    int int_X, int_T;

    if ( ( sscanf( (char*)argv[ 1], "%d", &int_X) != 1) || int_X <=0)
    {
        printf( "Wrong arguments\n");
        return 1;
    }
    if ( ( sscanf( (char*)argv[ 2], "%d", &int_T) != 1) || int_T <= 0)
    {
        printf( "Wrong arguments\n");
        return 1;
    }

    X = (real_t) int_X;
    T = (real_t) int_T;

	//// Init
    MPI_Init( &argc, &argv); 

    MPI_Comm_rank( MPI_COMM_WORLD, &comm_rank );
    MPI_Comm_size( MPI_COMM_WORLD, &comm_size );

    if ( comm_size < 1)
    {
        printf("comm_size must be bigger than 0");
        MPI_Finalize();
        return 0;
    }


    //// 
    real_t ** U;
    int * displs;
    int * counts;
    int N = (int)( ( real_t)int_X  / h);
    int i;

    displs = (int*)malloc( sizeof( int) * comm_size);
    counts = (int*)malloc( sizeof( int) * comm_size);

    optim_distrib( comm_size, N, displs, counts);


    //printf( "%d: from %d to %d ", comm_rank, displs[ comm_rank], displs[ comm_rank] + counts[ comm_rank]);

    U = (real_t **)malloc( sizeof( real_t**) * 2);

    counts[ comm_rank] += 2;
    U[ 0] = (real_t*)malloc( sizeof( real_t) * counts[ comm_rank]);
    U[ 1] = (real_t*)malloc( sizeof( real_t) * counts[ comm_rank]);

    double t_begin = MPI_Wtime();    

    // first iter
    for ( i = 0; i < counts[ comm_rank]; i++)
    {
        U[ 0][ i] = f_x( h * ( i + displs[ comm_rank]));
    } 

    int t = 0;
    while ( t * tau < T)
    {
        for ( i = 1; i < counts[ comm_rank] - 1; i++)
        {
            U[ 1][ i] = u( i, U);
        }

        if ( comm_rank % 2 == 0)
        {
            if ( comm_rank != comm_size - 1)
            {
                //printf  ( "%d: send + 1\n", comm_rank);
                MPI_Send ( &U[ 1][ counts[ comm_rank] - 2], 1, MPI_REAL, comm_rank + 1, 0, MPI_COMM_WORLD);
                                
                //printf  ( "%d: rcv  + 1\n", comm_rank);
                MPI_Recv ( &U[ 1][ counts[ comm_rank] - 1], 1, MPI_REAL, comm_rank + 1, 0, MPI_COMM_WORLD, status);

            }
            else
            {
                U[ 1][ counts[ comm_rank] -1] = right_u( i, U);
            }

            if ( comm_rank != 0)
            {
                //printf  ( "%d: send - 1\n", comm_rank);
                MPI_Send ( &U[ 1][ 1], 1, MPI_REAL, comm_rank - 1, 0, MPI_COMM_WORLD);
                
                //printf  ( "%d: rcv  - 1\n", comm_rank);
                MPI_Recv ( &U[ 1][ 0], 1, MPI_REAL, comm_rank - 1, 0, MPI_COMM_WORLD, status);
            }
            else
            {
                U[ 1][ 0] = f_t( t * tau);
            }
        }
        else
        {
            if ( comm_rank != comm_size - 1)
            {
                //printf  ( "%d: rcv  + 1\n", comm_rank);
                MPI_Recv ( &U[ 1][ counts[ comm_rank] -1], 1, MPI_REAL, comm_rank + 1, 0, MPI_COMM_WORLD, status);

                //printf  ( "%d: send + 1\n", comm_rank);
                MPI_Send ( &U[ 1][ counts[ comm_rank] -2], 1, MPI_REAL, comm_rank + 1, 0, MPI_COMM_WORLD);

            }
            else
            {
                U[ 1][ counts[ comm_rank] -1] = right_u( i, U);
            }

            //printf  ( "%d: rcv  - 1\n", comm_rank);
            MPI_Recv ( &U[ 1][ 0], 1, MPI_REAL, comm_rank - 1, 0, MPI_COMM_WORLD, status);

            //printf  ( "%d: send - 1\n", comm_rank);
            MPI_Send ( &U[ 1][ 1], 1, MPI_REAL, comm_rank - 1, 0, MPI_COMM_WORLD);            
        }

        real_t * tmp;
        tmp = U[ 1];
        U[ 1] = U[ 0];
        U[ 0] = tmp;

        t++;
    }

    double t_end = MPI_Wtime();

    if ( comm_rank == G_MAIN_RANK)
    {
        real_t * glob_U = (real_t*)malloc( sizeof( real_t) *  N);
        counts[ comm_rank] -= 2;

        MPI_Gatherv( U[ 1] + 1, counts[ comm_rank], MPI_REAL, glob_U, counts, displs, MPI_REAL, G_MAIN_RANK, MPI_COMM_WORLD);

        printf( "=======================================\nTime: %d\n==========================================\n", (int)( 1000000 * ( t_end - t_begin)));

        FILE * fh = fopen ( "res_Gvozdarev.txt", "w"); 
        for( i = 0; i < N; i++)
        {
            fprintf( fh, "%f\n", glob_U[ i]);
        }
        fclose( fh);
        free( glob_U);
    }
    else
    {
        counts[ comm_rank] -= 2;

        MPI_Gatherv( U[ 1] + 1, counts[ comm_rank], MPI_REAL, NULL, NULL, NULL, MPI_REAL, G_MAIN_RANK, MPI_COMM_WORLD);
    }


    free( displs);
    free( counts);

    free( U[ 0]);
    free( U[ 1]);
    free( U);

    MPI_Barrier( MPI_COMM_WORLD);
    MPI_Finalize();
    return 0;
}

real_t f_t( real_t t)
{
    return ( 0.5 * cos( 10.0 * t));
}

real_t f_x( real_t x)
{
    real_t x_2 = x * x;

    return ( exp( -0.5 * x_2) - 0.5 * exp( - 0.125 * x_2));
}

real_t u( int i, real_t ** U)
{
   return ( 0.5 * ( ( U[ 0][ i-1] + U[ 0][ i+1]) - tau / h * ( U[ 0][ i+1] - U[ 0][ i-1])));
}

real_t right_u( int i, real_t ** U)
{
       return ( U[ 0][ i] - tau / h * ( U[ 0][ i-1] - U[ 0][ i]));
}

void optim_distrib( int comm_size, int data_size, int * displs, int * counts)
{
    int to_all = data_size / comm_size;
    int mod = data_size % comm_size;

    int i;
    int curr = 0;
    for ( i = 0; i < mod; i++)
    {
        counts[ i] = to_all + 1;
        displs[ i] = curr;
        curr += counts[ i];
    }
    for ( i = mod; i < comm_size; i++)
    {
        counts[ i] = to_all;
        displs[ i] = curr;
        curr += counts[ i];
    }
}

